
from libatomism import *
from math import *
import array 
import matplotlib.pyplot as plt
"""
This script shows how to model a molecule, the hydrogen molecule is taken as example. 
The following feature are described:
    - managment of units
    - z-matrix generalized coordinates
    - equation of motion: kinetic operator and potential energy surface

We recomand you to use the interpreter ipython for the numerous sugar features it offers.

Some tips:
    - all objects can be printed by 'print myobjetc'
    - to obtain the documentation of a class, input 'myClass?' and press enter (ipython only)
    - to obtain the list of method in an object, input 'myobject.' and press tab (ipython only)
    - to obtain the doxumentation on a method in an object, input 'myobject.myMethod?' and press enter (ipython only)
"""


""" Set the logger to print only error messages
"""
msLogger.setPriorityMin(Priority.ERROR)

""" Create a units manager object, it will be very often used to instantiate 
the different classes: it defines the unit of the intput and output of the objects
 by default. A units manager object is also usefull for conversion purposes 
(we'll see that later).
"""
unit = msUnitsManager.New("Angstrom Degree amu kcal/mol ps")

""" Let's define the H2 ZMatrix:
H
H 0 0.74191
The z-matrix is the object who define the link between coordinates value
and position of the elements. If you want to define a custom problem 
(e.g. a chain of pendulums), you only have to write a similar class and
 you should look at the msZMat source files for documentation.  
"""
zmat = msZMat.New(unit).setId("zmat").set("H H 0 0.74191")
"""
When you print an object, the all class hierarchy of the object is described.
At the first level (i.e. base class), you have the description of the 'database mapper':
its purpose is to translate the c++ class structure into tree structure in order to
be saved or loaded. 
Under it, the first derived class from the tree mapper, the description of the 
'Physical interface' interface. The purpose of this interface is to define a unique 
framework to define units and parameters.
At the third level you have the description of the entity class: a collection 
of elements with their position and masse.
Finally the description of the Z-matrix is providen.
You can obtain more information about all classes in their doxygen documentation.
""" 
print zmat

""" Now a system is defined. A system is just a collection of entity (parent class
of Z-Matrix), that you add to the container:
"""
system = msSystem.New(unit).addEntity(zmat)
print system

""" If you have VTK and if you linked atomism with it, you can display the system
by uncommenting the two next lines
"""
# viewer = msVisualizerAtoms.New()
# viewer.watch(system)

"""
At this point the system is defined: there is some elements (atoms here) associated
to a mass and their position is defined from some coordinates (Z-Matrix here). 
The next step is the definition of motion: how the elements are moving w/ some 
generalized coordinates, and what is the energy cost of it.
Lets consider the vibration H-H in H2.
step 1: definition of the generalized coordinates, it contains one variable,
 the separation H-H.
"""
gcoors = msGeneralizedCoordinates.New(unit)
q0     = msScalarVariable.New("Angstrom").set(0.7414, 0.1, 5, 0.01, 0.1).setId("q0")

""" try 'q0.set?' for the description of the method. 
The generalized coordinates object is a container for scalar variable: 
""" 
gcoors.addVar(q0)
print gcoors

""" step 2: definition of the potential energy.
The potential energy can be any object deriving from scalar functions,
here we are using a parser type to define the Morse potential used by Martson 
 and Bali-Kunti in their study. 
see C. Clay Martson and G. G. Balint-Kurti, J Chem Phys 91(6): 3571 (1989)
"""
epot = msParser.New( unit.getEnergyStr() ).setCoordinates(gcoors)
epot.setExpression("De * ( 1 - exp( - beta*( q0 - re ) ))^2").setId("f=Morse")

""" The unit object allow easy conversion of the tabulated parameters to 
our unit system:
"""
De   = unit.convert("eV",4.7457)
beta = unit.convert("m^-1",1.94196e10)
re   = unit.convert("m",0.74191e-10)

epot.setConstant("De",De).setConstant("beta",beta).setConstant("re",re)
print epot

""" step 3: definition of the kinetic operator. 
The kinetic operator describes how the masses in the 'system' object move with respect 
to the generalized coordinates.
"""
kinop = msKineticOperator.New(unit).set(gcoors,system)

""" The kinetic operator is a container of dynamic Degrees Of Freedom (DOF);
these DOF link a coordinate in an entity to a scalar function.
Here, we want to link the separation H-H in the z-matrix (DOF number 0 of 'zmat') 
by the scalar function f(q0)=q0.
"""
kinop.addDynamicDof( zmat, 0, msParser.New("Angstrom").setCoordinates(gcoors).setExpression("q0").setId("f=q0") )
print kinop

""" PES and kinetic operator defined a motion, in the case of 1D motion, a
quantum resolution of the equation of motion is available.
""" 
qmotion = msQuantumMotion1d.New(unit).setEqOfMotion(kinop,epot)
print qmotion

""" In the previous ouput, you can look at the default value of the parameters.
To change the value of a parameter: 
"""
qmotion.getParameters().setParameter("BasisSize",501,unit)
qmotion.getParameters().setParameter("Emax",De,unit)

""" All objects derived from the msMotion class contains several themodynamic and DOS related
functions; see the documentation of the class msMotion. For instance:
""" 
Q   = qmotion.Q(1000)
S   = qmotion.S(1000)
Cp  = qmotion.Cp(1000)
dos = qmotion.DOS(10)
sos = qmotion.SOS(10)
""" We want to compare the results with those of an harmonic oscilator.
From the PES 'epot', we define its harmonic approximation at the 
equilibrium separation.
""" 
epotHO = msParser.New(unit.getEnergyStr()).setCoordinates(gcoors).setExpression("alpha * ( q0 - re )^2").setId("f=HO")

gcoors[0].Value = re

""" By default scalar functions implement a numeric derivation procedure.
The indice given in the function call is the indice of the variable we are deriving from.
"""
epotHO.setConstant("alpha", epot.partialSecondDerivative(0) ).setConstant("re",re)

oscillsHO = msQuantumMotion1d.New(unit).setEqOfMotion(kinop,epotHO)
oscillsHO.getParameters().setParameter("BasisSize",501,"")
oscillsHO.getParameters().setParameter("Emax",De,unit)

""" Now we need to define overall rotational and overall translational motions.
These modes of motion are not studied bt their equation of motion, but by 
analytically derived statistical expressions.
Be carefull at the geometry of your system at the time your are fixing the 
moment of inertia. Here, we recompute the atomic position at the current Value 
of q0 (i.e. re, see 4 lines ahead)
"""
kinop.setDynamicDofs()
system.computeCartCoordinates()
rotor = msLinearRigidRotor.New(unit).setMmtOfInertia(system).setSymmetryNumber(2)
rotor.getParameters().setParameter("Emax",unit.convert("kcal/mol",200),unit)

trans = msTranslation.New(unit).setSystem(system)

""" At this point everything is defined. Let's realized some computation
and plot the results.
"""
qmotion.computeEigenValues()
oscillsHO.computeEigenValues()

""" In this particular case the kinetic function (or reduced mass because q0 is a length)
 is constant w/q0. That means that the kinetic function can be freezed to its current 
value.
"""
qmotion.freezeKinValue( kinop.kineticFunction() )
oscillsHO.freezeKinValue( kinop.kineticFunction() )

""" get the number 0, 4 and 9 eigenvectors, they will be plot later 
"""
ev0  = qmotion.getEigenvector(0)
ev5  = qmotion.getEigenvector(4)
ev10 = qmotion.getEigenvector(9)
ev0HO = oscillsHO.getEigenvector(0)
ev5HO  = oscillsHO.getEigenvector(4)
ev10HO = oscillsHO.getEigenvector(9)


rotor.computeEigenValues()
""" Occupency probability of the rotor eigenmodes at 500 and 2000 K
"""
Jdist500=rotor.getOccupencyProbability(500)
Jdist2000=rotor.getOccupencyProbability(2000)

""" Compute the total entropy for the Morse and Harmonic PES 
"""
s=[]
sHO=[]
Tlist  = [200, 300, 500, 800, 1200, 1500, 2000, 2500, 3000] 
for T in Tlist: 
    s.append(   trans.S(T)  + rotor.S(T) + qmotion.S(T)   + unit.convert("J/mol",R) )
    sHO.append( trans.S(T)  + rotor.S(T) + oscillsHO.S(T) + unit.convert("J/mol",R) )

""" Same for calorific capacity
"""
cp=[]
cpHO=[]
for T in Tlist:
    cp.append(   trans.Cv(T) + rotor.Cv(T) + qmotion.Cv(T)   +  unit.convert("J/mol",R) )
    cpHO.append( trans.Cv(T) + rotor.Cv(T) + oscillsHO.Cv(T) +  unit.convert("J/mol",R) )

    
""" See the documentation of matplotlib for the plotting part.
All the object derived from a scalar function can return a python list to 
be plotted using the 'getXSample1d' and 'getYSample1d' methods. 
In these functions the first variable ie the indice of the coordinates 
along witch you want to sample, and the second the number of sample you want.
"""
fig, axes = plt.subplots(3, 2)

""" The following four lines define the analytical eigenvalues of the Morse potential 
in the H2 molecule. see C. Clay Martson and G. G. Balint-Kurti, J Chem Phys 91(6): 3571 (1989)
"""
scoor=[q0.Min,q0.Max]
y0_analytic=[unit.convert("hartree",9.86922e-3),unit.convert("hartree",  9.86922e-3)]
y4_analytic=[unit.convert("hartree",7.847518e-2),unit.convert("hartree", 7.847518e-2)]
y9_analytic=[unit.convert("hartree",1.3836303e-1),unit.convert("hartree",1.3836303e-1)]

"""on top right subplot, we plot the potential energy and the eigenvectors with 
 a sampling of 500 along the first coordinates of 'gencoors' (indiced 0, i.e. q0);
 the number '600' is a scaling factor for having a nice drawing of the eigenvectors
"""
axes[0, 0].plot(    epotHO.getXSample1d(0,500), epotHO.getYSample1d(0,500), "gray", 
	            ev0HO.getXSample1d(0,500) , [ 600*x + oscillsHO.eigenvalue(0)  for x in ev0HO.getYSample1d(0,500)] ,"gray",  
	            ev5HO.getXSample1d(0,500) , [ 600*x + oscillsHO.eigenvalue(4)  for x in ev5HO.getYSample1d(0,500)] ,"gray",
	            ev10HO.getXSample1d(0,500), [ 600*x + oscillsHO.eigenvalue(9) for x in ev10HO.getYSample1d(0,500)],"gray",
		    epot.getXSample1d(0,500), epot.getYSample1d(0,500), "blue", 
	            ev0.getXSample1d(0,500) , [ 600*x + qmotion.eigenvalue(0)  for x in ev0.getYSample1d(0,500)] ,"blue",  
	            ev5.getXSample1d(0,500) , [ 600*x + qmotion.eigenvalue(4)  for x in ev5.getYSample1d(0,500)] ,"blue",
	            ev10.getXSample1d(0,500), [ 600*x + qmotion.eigenvalue(9) for x in ev10.getYSample1d(0,500)],"blue",
		    scoor,y0_analytic,"red",
		    scoor,y4_analytic,"red",
	            scoor,y9_analytic,"red",
		    )
	            
axes[0, 0].axis((0,3.5,0,unit.convert("kcal/mol",160)))
axes[0, 0].set_xlabel("H-H ["+unit.getLengthStr()+"]",fontsize=20)
axes[0, 0].set_ylabel("energy ["+unit.getEnergyStr()+"]",fontsize=20)
axes[0, 0].tick_params(labelsize=20)

# On the top left subplot we plot the sum of states of 'qmotion' and compare it 
# with the one of the harmonic oscillator:
axes[0, 1].plot( qmotion.SOS().getXSample1d(0,500),    qmotion.SOS().getYSample1d(0,500), "-",
                 oscillsHO.SOS().getXSample1d(0,500),  oscillsHO.SOS().getYSample1d(0,500), "gray")
axes[0, 1].set_xlabel("vibrational energy ["+unit.getEnergyStr()+"]",fontsize=20)
axes[0, 1].set_ylabel("vibrational sum of states",fontsize=20)
axes[0, 1].tick_params(labelsize=20)
	
# On the middle left subplot we plot the occupency probabilities of the rotor w/
# to the eigenvalues:
axes[1, 1].plot( Jdist500.getXSample1d(0,1000),   Jdist500.getYSample1d(0,1000),  "-",
                 Jdist2000.getXSample1d(0,1000),  Jdist2000.getYSample1d(0,1000), "-")
axes[1, 1].set_xlabel("angular momentum number",fontsize=20)
axes[1, 1].set_ylabel("Occupency probability",fontsize=20)
axes[1, 1].tick_params(labelsize=20)

# On the middle right sublot we plot the rotor Density Of States and Sum Of States,
# We need in a high resolution to represent the diracs of the DOS
axes[1, 0].plot(rotor.DOS().getXSample1d(0,10000),  rotor.DOS().getYSample1d(0,10000), "-",
                rotor.SOS().getXSample1d(0,1000),   rotor.SOS().getYSample1d(0,1000), "-")  
axes[1, 0].set_xlabel("rotational energy ["+unit.getEnergyStr()+"]",fontsize=20)
axes[1, 0].set_ylabel("rotational SOS and DOS",fontsize=20)
axes[1, 0].tick_params(labelsize=20)
axes[1, 0].axis((0,unit.convert("kcal/mol",200),0,1400))
# Now we compute the entropy and calorific capacity lists and plot them in a subplots.
# Again, the results are compared to those of the harmonic oscillator.

axes[2, 0].plot( Tlist, s , "-", Tlist, sHO)
axes[2, 0].set_xlabel("Temperature ["+unit.getTempStr()+"]",fontsize=20)
axes[2, 0].set_ylabel("Total entropy ["+unit.getEnergyStr()+"]",fontsize=20)
axes[2, 0].tick_params(labelsize=20)

axes[2, 1].plot( Tlist, cp, "-", Tlist, cpHO,'gray')
axes[2, 1].set_xlabel("Temperature ["+unit.getTempStr()+"]",fontsize=20)
axes[2, 1].set_ylabel("Total Cp ["+unit.getEnergyStr()+"]",fontsize=20)
axes[2, 1].tick_params(labelsize=20)

# this is the experimental values, coming from the janaf table,  
# see http://kinetics.nist.gov/janaf/html/H-050.html, defined in J/mol/K
TListexp  =  [200,     300,     500,     800,     1200,   1500,    2000,    2500,    3000   ] 
ExpListS   = [119.412, 130.858, 145.737, 159.548, 171.79, 178.846, 188.418, 196.243, 202.891]   
ExpListCp =  [27.447,  28.849,  29.260,  29.624,  30.992, 32.298,  34.280,  35.842,  37.087 ] 
 
axes[2, 0].plot( TListexp, [ unit.convert("J/mol",x) for x in ExpListS],   "r*" ,markersize=15)
axes[2, 1].plot( TListexp, [ unit.convert("J/mol",x) for x in ExpListCp],  "r*" ,markersize=15)

fig.set_size_inches(20,20)
fig.savefig('H2_0.png')
fig.show()


 
